import threading
import heapq
import collections
from queue import (Queue, PriorityQueue)
from action_msgs.msg import GoalStatus
from rclpy.action import (ActionServer, ActionClient, CancelResponse, GoalResponse)
from rclpy.callback_groups import ReentrantCallbackGroup
from rclpy.node import Node
from semantic_nav.utils import WrappedComparable


class ActionServerTemplate(Node):
    """ Base template for creating an action server. """
    def __init__(self, action_type, action_name, namespace='', execute_func=None):
        super().__init__(node_name=action_name + '_server', namespace=namespace)
        self.action_type = action_type
        self._callback_group = ReentrantCallbackGroup()
        if execute_func is None:
            self.execute = lambda node, goal_handle: action_type.Result()
        else:
            self.execute = execute_func

        self._server = ActionServer(
            self,
            action_type,
            action_name,
            handle_accepted_callback=self.handle_accepted_callback,
            execute_callback=self.execute_callback,
            goal_callback=self.goal_callback,
            cancel_callback=self.cancel_callback,
            callback_group=self._callback_group
        )

    def __del__(self):
        self._server.destroy()
        super().destroy_node()

    def goal_callback(self, goal_request):
        """Accept or reject a client request to begin an action."""
        self.get_logger().info(f'Received goal request: {goal_request}')
        return GoalResponse.ACCEPT

    def handle_accepted_callback(self, goal_handle):
        goal_handle.execute()

    def cancel_callback(self, goal_handle):
        """Accept or reject a client request to cancel an action."""
        self.get_logger().info('Received cancel request')
        return CancelResponse.ACCEPT

    def is_interrupted(self, goal_handle) -> bool:
        """ Check if the goal is interrupted, needed to be called in each cycle of execute"""
        if goal_handle.is_cancel_requested:
            # if so, we cancel the goal handle
            goal_handle.canceled()
            self.get_logger().info('Goal canceled')
            return True
        if not goal_handle.is_active:
            self.get_logger().info('Goal aborted')
            return True
        return False

    async def execute_callback(self, goal_handle):
        """Execute a goal."""
        return self.execute(self, goal_handle)


class SingleRequestServer(ActionServerTemplate):
    """ action server which only one request can be processed at the same time. """
    def __init__(self, action_type, action_name, namespace='', execute_func=None):
        super().__init__(action_type, action_name, namespace=namespace, execute_func=execute_func)
        self._goal_lock = threading.Lock()
        self._current_goal_handle = None

    def handle_accepted_callback(self, goal_handle):
        """Start or defer execution of an already accepted goal."""
        with self._goal_lock:
            if self._current_goal_handle is not None and self._current_goal_handle.is_active:
                self.get_logger().info('Aborting previous goal')
                self._current_goal_handle.abort()

            self._current_goal_handle = goal_handle

        goal_handle.execute()

    def is_interrupted(self, goal_handle) -> bool:
        with self._goal_lock:
            return super().is_interrupted(goal_handle)

    def is_succeeded(self, goal_handle) -> bool:
        """Call this function when you want to set the goal handle to success"""
        with self._goal_lock:
            if not super().is_interrupted(goal_handle):
                goal_handle.succeed()
                return True
        return False

    def execute_callback(self, goal_handle):
        """Execute a goal."""
        return self.execute(self, goal_handle)


class FifoQueServer(ActionServerTemplate):
    """ action server with first-in-first-out queue. """
    def __init__(self, action_type, action_name, namespace='', execute_func=None):
        super().__init__(action_type, action_name, namespace=namespace, execute_func=execute_func)
        self._goal_queue = collections.deque()  # The first-in-first-out queue stores the goals to be processed
        self._goal_lock = threading.Lock()
        self._current_goal_handle = None

    def clear(self):
        with self._goal_lock:
            self._goal_queue.clear()
            self._current_goal_handle = None

    def handle_accepted_callback(self, goal_handle):
        """Start or defer execution of an already accepted goal."""
        with self._goal_lock:
            if self._current_goal_handle is not None:
                self._goal_queue.append(goal_handle)
            else:
                # Start goal execution right away
                self._current_goal_handle = goal_handle
                self._current_goal_handle.execute()

    def execute_callback(self, goal_handle):
        """Execute a goal."""
        try:
            return self.execute(self, goal_handle)
        finally:
            with self._goal_lock:
                try:
                    # Start execution of the next goal in the queue.
                    self._current_goal = self._goal_queue.popleft()
                    self.get_logger().info('Next goal pulled from the queue')
                    self._current_goal.execute()
                except IndexError:
                    # No goal in the queue.
                    self._current_goal = None


class PriorQueServer(SingleRequestServer):
    """ action server with priority queue. """
    def __init__(self, action_type, action_name, namespace='', execute_func=None):
        if not hasattr(action_type.Goal(), 'priority'):
            raise TypeError('Action type must have a priority field')
        super().__init__(action_type, action_name, namespace=namespace, execute_func=execute_func)
        self._goal_queue = []  # The priority queue stores the goals to be processed
        self._count = int(0)
        self.aborted = bool(False)

    def handle_accepted_callback(self, goal_handle):
        """Start or defer execution of an already accepted goal."""
        with self._goal_lock:
            if self._current_goal_handle is not None and self._current_goal_handle.is_active:
                if self._current_goal_handle.request.priority > goal_handle.request.priority:
                    self.get_logger().info('Aborting previous goal')
                    # Abort the existing goal
                    self._current_goal_handle.abort()
                    self.aborted = True
                else:
                    # Put incoming goal in the queue
                    self.get_logger().info('Goal put in the queue')
                    # Priority queue has the feature of first-in-last-out with the same priority
                    # Add sequence counting to make sure the goal with same priority is executed in the order they are received
                    heapq.heappush(self._goal_queue, WrappedComparable((goal_handle.request.priority, self._count), goal_handle))
                    self._count += 1
                    return

            # Start goal execution right away
            self._current_goal_handle = goal_handle
            self._current_goal_handle.execute()

    def execute_callback(self, goal_handle):
        """Execute a goal."""
        try:
            return self.execute(self, goal_handle)
        finally:
            with self._goal_lock:
                if not self.aborted:
                    if self._goal_queue:
                        # Start execution of the next goal in the queue.
                        self._current_goal_handle = heapq.heappop(self._goal_queue).value
                        self.get_logger().info('Next goal pulled from the queue')
                        self._current_goal_handle.execute()
                    else:
                        # No goal in the queue.
                        self._count = int(0)
                        self._current_goal_handle = None
                else:
                    self.aborted = False


class ActionClientTemplate(Node):
    """ Base template for creating an action client. """
    def __init__(self, action_type, action_name, namespace='', func_succeeded=None, func_failed=None, func_feedback=None):
        super().__init__(node_name=action_name + '_client', namespace=namespace)
        self.action_type = action_type
        self._client = ActionClient(self, action_type, 'semantic_nav')
        self._goal_handle = None
        if func_succeeded is None:
            self.execute_succeeded = lambda result: self.get_logger().info(f"Execution succeeded: {result}")
        else:
            self.execute_succeeded = func_succeeded
        if func_failed is None:
            self.execute_failed = lambda result: self.get_logger().info(f"Execution failed: {result}")
        else:
            self.execute_failed = func_failed
        if func_feedback is None:
            self.execute_feedback = lambda feedback: self.get_logger().info(f"Received feedback: {feedback}")
        else:
            self.execute_feedback = func_feedback

    def __del__(self):
        self._client.destroy()
        super().destroy_node()

    def send_goal(self, goal):
        self.get_logger().info("Waiting for action server...")
        self._client.wait_for_server()
        self.get_logger().info(f"Send goal: {goal}")
        self._send_goal_future = self._client.send_goal_async(goal, feedback_callback=self.feedback_callback)
        self._send_goal_future.add_done_callback(self.goal_response_callback)

    def cancel_goal(self):
        if self._goal_handle is not None:
            self.get_logger().info("Cancel goal")
            cancel_future = self._goal_handle.cancel_goal_async()
            cancel_future.add_done_callback(self.cancel_done)
            return True
        self.get_logger().info("No active goal to cancel")
        return False

    def cancel_done(self, future):
        cancel_response = future.result()
        if len(cancel_response.goals_canceling) > 0:
            self.get_logger().info("Goal successfully canceled")
        else:
            self.get_logger().info("Goal failed to cancel")

    def goal_response_callback(self, future):
        goal_handle = future.result()
        if not goal_handle.accepted:
            self.get_logger().info("Goal rejected")
            return

        self._goal_handle = goal_handle
        self.get_logger().info("Goal accepted")
        self._get_result_future = goal_handle.get_result_async()
        self._get_result_future.add_done_callback(self.get_result_callback)

    def feedback_callback(self, feedback_msg):
        # feedback_msg:
        # - goal     (unique_identifier_msgs.msg.UUID)
        # - feedback (action_type.Feedback)
        self.execute_feedback(feedback_msg.feedback)

    def get_result_callback(self, future):
        result_msg = future.result()
        # result_msg:
        # - status
        # - result (action_type.Result)
        if result_msg.status == GoalStatus.STATUS_SUCCEEDED:
            self.execute_succeeded(result_msg.result)
        else:
            self.execute_failed(result_msg.result)


class FifoClient(ActionClientTemplate):
    """ action client with first-in-first-out queue. """
    def __init__(self, action_type, action_name, namespace='', func_succeeded=None, func_failed=None, func_feedback=None):
        super().__init__(
            action_type=action_type,
            action_name=action_name,
            namespace=namespace,
            func_succeeded=func_succeeded,
            func_failed=func_failed,
            func_feedback=func_feedback
        )
        self._busy_lock = threading.Lock()
        self._busy: bool = False
        self._wait_set = Queue()

    def __del__(self):
        super().__del__()
        self._wait_set = None

    def send_goals(self, goals):
        for goal in goals:
            self._wait_set.put(goal)
        if self.get_busy():
            return
        if not self._wait_set.empty():
            self.send_goal(self._wait_set.get())

    def cancel_goals(self):
        # clear the waiting goals queue
        while not self._wait_set.empty():
            self._wait_set.get()
        return self.cancel_goal()

    def get_result_callback(self, future):
        super().get_result_callback(future)
        if self._wait_set.empty():
            self.set_busy(False)
            return
        goal = self._wait_set.get()
        self.send_goal(goal)

    def set_busy(self, busy: bool):
        with self._busy_lock:
            self._busy = busy

    def get_busy(self):
        with self._busy_lock:
            return self._busy


class PriorQueClient(FifoClient):
    """ action client with priority queue. """
    def __init__(self, action_type, action_name, namespace='', func_succeeded=None, func_failed=None, func_feedback=None):
        if not hasattr(action_type.Goal(), 'priority'):
            raise TypeError('Action type must have a priority field')
        super().__init__(
            action_type=action_type,
            action_name=action_name,
            namespace=namespace,
            func_succeeded=func_succeeded,
            func_failed=func_failed,
            func_feedback=func_feedback
        )
        self._success_lock = threading.Lock()
        self._success: bool = False
        self._goal_queue = PriorityQueue()

    def __del__(self):
        super().__del__()
        self._goal_queue = None

    def send_goals(self, goals):
        if self.get_busy():
            self._wait_set.put(goals)
            return
        self._send_goals(goals)

    def _send_goals(self, goals):
        self.set_busy(True)
        self.set_result(False)
        count = int(0)
        for goal in goals:
            # Priority queue has the feature of first-in-last-out with the same priority
            # Add sequence counting to make sure the goal with same priority is executed in the order they are received
            self._goal_queue.put(WrappedComparable((goal.priority, count), goal))
            count += 1
        if not self._goal_queue.empty():
            self.send_goal(self._goal_queue.get().value)

    def cancel_goals(self):
        # clear the waiting goals queue
        while not self._goal_queue.empty():
            self._goal_queue.get()
        return super().cancel_goals()

    def get_result_callback(self, future):
        super(FifoClient, self).get_result_callback(future)

        if not self._goal_queue.empty():
            self.send_goal(self._goal_queue.get().value)
            return

        if self._wait_set.empty():
            self.set_busy(False)
            return
        goals = self._wait_set.get()
        self._send_goals(goals)

    def set_result(self, success: bool):
        with self._success_lock:
            self._success = success

    def get_result(self):
        with self._success_lock:
            return self._success
